package codegen


import . "compiler/ast"
import . "compiler/lexer"
import . "vm"
// import "fmt"


var arithAndBitwiseBinops = map[int] int {
	TOKEN_OP_ADD:	OP_ADD,
	TOKEN_OP_SUB:	OP_SUB,
	TOKEN_OP_MUL:	OP_MUL,
	TOKEN_OP_MOD:	OP_MOD,
	TOKEN_OP_POW:	OP_POW,
	TOKEN_OP_DIV:	OP_DIV,
	TOKEN_OP_IDIV:	OP_IDIV,
	TOKEN_OP_BAND:	OP_BAND,
	TOKEN_OP_BOR:	OP_BOR,
	TOKEN_OP_BXOR:	OP_BXOR,
	TOKEN_OP_SHL:	OP_SHL,
	TOKEN_OP_SHR:	OP_SHR,
}


// Lua采用词法作用域
// 在函数内部，某个局部变量的作用域时包围该变量的最内层语句块
// 简单来说，每个块都会制造一个新的作用域
// 在块的内部可以使用局部变量声明语句声明(并初始化)局部变量
// 每个局部变量都会占用一个寄存器索引
// 当块结束后，作用域也随之消失，局部变量不复存在，占用的寄存器也会被回收
type locVarInfo struct {

	prev		*locVarInfo		// 使locVarInfo结构体成为单向链表节点
	name 		string			// 记录局部变量名
	scopeLv		int 			// 记录局部变量所在的作用域层次
	slot		int 			// 记录与局部变量名绑定的寄存器索引
	captured	bool 			// 记录局部变量是否被闭包捕获

}


// Upvalue实际上就是闭包按照词法作用域捕获的外围函数中的局部变量
// 和局部变量类似，我们也需要把Upvalue名和外围函数的局部变量绑定
// 不过由于Upvalue名仅能绑定唯一的Upvalue，所以不需要使用链表结构
type upvalInfo struct {

	// 如果Upvalue捕获的是直接外围函数的局部变量
	// 则locVarSlot字段记录该局部变量所占的寄存器索引
	// 否则Upvalue已经被直接外围函数捕获
	// upvalueIndex字段记录该Upvalue在直接外围函数Upvalue表中的索引
	locVarSlot		int 
	upvalueIndex	int 

	// 记录Upvalue在函数中出现的顺序
	index 			int 

}


type funcInfo struct {

	parent			*funcInfo			// 使我们能够定位到外围函数的局部变量表和Upvalue表

	// 常量表，存放函数体出现的nil、布尔、数字或者字符串字面量
	// 键是常量值，值是常量在表中的索引
	constants	map[interface{}]int

	// Lua虚拟机是基于寄存器的机器，因此我们在生成指令时需要进行寄存器分配
	// 简单来说，我们需要给每一个局部变量和临时变量都分配一个寄存器
	// 在局部变量退出作用域或者临时变量使用完毕之后，回收寄存器
	usedRegs	int 		// 已分配的寄存器数量
	maxRegs		int 		// 需要的最大寄存器数量

	scopeLv 	int 		// 记录当前作用域层次
	locVars		[]*locVarInfo		// 按顺序记录函数内部声明的全部局部变量
	// 记录当前生效的局部变量
	// 作用域层次从0开始，每进入一个作用域就加1
	locNames		map[string]*locVarInfo

	// 使用数组来记录循环块内待处理的跳转指令
	// 数组长度和块的深度对应，通过判断数组元素(也是数组)是否为nil
	// 可以判断块是否是循环块
	breaks			[][]int

	upvalues		map[string]upvalInfo		// 存放Upvalue表

	insts			[]uint32			// 存储编码后的指令

	// 函数原型是递归结构，因为其内部还可以有子函数原型
	// 与此相应，funcInfo结构体也需要是递归结构
	subFuncs		[]*funcInfo		// 存放子函数信息
	numParams		int 	// 参数数量
	isVararg		bool 	// 是否是可变参数

}


// 创建funcInfo结构体实例
func newFuncInfo(parent *funcInfo, fd *FuncDefExp) *funcInfo {

	return &funcInfo{
		parent:		parent,
		subFuncs:	[]*funcInfo{},
		constants:	map[interface{}]int{},
		upvalues:	map[string]upvalInfo{},
		locNames:	map[string]*locVarInfo{},
		locVars:	make([]*locVarInfo, 0, 8),
		breaks:		make([][]int, 1),  // 默认长度为1，这是因为主函数本身也是一个作用域.
		insts:		make([]uint32, 0, 8),
		numParams:	len(fd.ParList), 
		isVararg:	fd.IsVararg,
	}

}


// 返回常量在表中的索引
// 如果常量不在表里，先把常量表放入表里，再返回索引
// 索引从0开始递增
func (self *funcInfo) indexOfConstant(k interface{}) int {

	if idx, found := self.constants[k]; found {

		return idx
	}

	idx := len(self.constants)
	self.constants[k] = idx

	return idx

}


// 分配一个寄存器，必要时更新最大寄存器数量
// 并返回寄存器索引
// 寄存器索引是从0开始的，且不能超过255
func (self *funcInfo) allocReg() int {

	self.usedRegs++

	if self.usedRegs >= 255 {
		panic("function or expression needs too many registers")
	}

	if self.usedRegs > self.maxRegs {
		self.maxRegs = self.usedRegs
	}

	return self.usedRegs - 1

}


// 回收最近分配的寄存器
func (self *funcInfo) freeReg() {
	if self.usedRegs <= 0 {
		panic("usedRegs <= 0 !")
	}
	self.usedRegs--

}


// 分配连续的n个寄存器，返回第一个寄存器的索引
func (self *funcInfo) allocRegs(n int) int {
	if n <= 0 {
		panic("n <= 0 !")
	}
	for i := 0; i < n; i++ {

		self.allocReg()

	}

	return self.usedRegs - n

}


// 回收最近分配的n个寄存器
func (self *funcInfo) freeRegs(n int) {
	if n < 0 {
		panic("n < 0 !")
	}
	for i := 0; i < n; i++ {

		self.freeReg()

	}

}


// 进入新的作用域
// 对进入的作用域是否属于循环块进行标记
func (self *funcInfo) enterScope(breakable bool) {

	self.scopeLv++

	if breakable {
		self.breaks = append(self.breaks, []int{})		// 循环块
	} else {
		self.breaks = append(self.breaks, nil)			// 非循环块
	}

}


// 退出作用域
// 在退出时修复跳转指令
func (self *funcInfo) exitScope() {

	pendingBreakJmps := self.breaks[len(self.breaks)-1]
	self.breaks = self.breaks[:len(self.breaks)-1]

	a := self.getJmpArgA()
	for _, pc := range pendingBreakJmps {
		sBx := self.pc() - pc
		i := (sBx + MAXARG_sBx)<<14 | a<<6 | OP_JMP
		self.insts[pc] = uint32(i)
	}

	self.scopeLv--

	for _, locVar := range self.locNames {

		if locVar.scopeLv > self.scopeLv {	// 离开作用域

			self.removeLocVar(locVar)
		}
	}

}


// 退出作用域后，需要删除该作用域内的局部变量(解绑局部变量，回收寄存器)
// 首先回收寄存器，然后看是否有其它同名局部变量
// 如果无，则直接解绑局部变量名即可
// 如果有，且在同一个作用域内，则递归调用removeLocVar()方法进行处理
// 否则，同名局部变量在更外层的作用域里，我们需要把局部变量名与该局部变量重新绑定
func (self *funcInfo) removeLocVar(locVar *locVarInfo) {

	self.freeReg()		// 回收寄存器

	if locVar.prev == nil {

		delete(self.locNames, locVar.name)
	} else if locVar.prev.scopeLv == locVar.scopeLv {

		self.removeLocVar(locVar.prev)
	} else {

		self.locNames[locVar.name] = locVar.prev
	}

}


// 在当前作用域添加一个局部变量
// 返回其分配的寄存器索引
func (self *funcInfo) addLocVar(name string) int {

	newVar := &locVarInfo {
		name:	name,
		prev:	self.locNames[name],
		scopeLv:	self.scopeLv,
		slot:	self.allocReg(),		// 分配寄存器
	}

	self.locVars = append(self.locVars, newVar)
	self.locNames[name] = newVar

	return newVar.slot

}


// 检查局部变量名是否已经和某个寄存器绑定
// 如果是则返回寄存器索引
// 否则返回-1
func (self *funcInfo) slotOfLocVar(name string) int {

	if locVar, found := self.locNames[name]; found {

		return locVar.slot
	}

	return -1

}


// 把break语句对应的跳转指令添加到最近的循环块里
// 如果找不到循环块则调用panic()函数汇报错误
func (self *funcInfo) addBreakJump(pc int) {

	for i := self.scopeLv; i >= 0; i-- {		// 第0块不为空
		if self.breaks[i] != nil {  // 循环块
			self.breaks[i] = append(self.breaks[i], pc)
			return
		}
	}

	panic("<break> at line ? not inside a loop!")

}


// 判断名字是否已经和Upvalue绑定，如果是，返回Upvalue索引
// 否则，尝试绑定，然后返回索引
// 如果绑定失败，则返回-1
func (self *funcInfo) indexOfUpval(name string) int {

	if upval, ok := self.upvalues[name]; ok {

		return upval.index
	}

	if self.parent != nil {

		if locVar, found := self.parent.locNames[name]; found {

			idx := len(self.upvalues)
			self.upvalues[name] = upvalInfo{locVar.slot, -1, idx}

			locVar.captured = true

			return idx
		}

		if uvIdx := self.parent.indexOfUpval(name); uvIdx >= 0 {

			idx := len(self.upvalues)
			// 同一个变量在多个作用域内可能会有多个upvalue
			// 或者说每一个函数针对同一个被捕获的局部变量都会存入自身的upvalue列表
			self.upvalues[name] = upvalInfo{-1, uvIdx, idx}  

			return idx
		}
	}

	return -1

}


// 当局部变量退出作用域时，我们调用closeOpenUpvals()方法将处于开启状态的Upvalue闭合
// 如有需要处理的局部变量，该方法将会产生一条JMP指令
// 其操作数A给出需要处理的第一个局部变量的寄存器索引
func (self *funcInfo) closeOpenUpvals() {

	a := self.getJmpArgA()
	if a > 0 {
		self.emitJmp(a, 0)
	}
}


// 获取JMP指令的操作数A
// Jmp指令并不会精准地去闭合该作用域内的某个upvalue，而是采用一刀切的方式
// 如果该作用域内没有变量被捕获，那么就不去闭合upvalue(实际上也没有), 指令a操作数的值为0
// 如果该作用域内有局部变量被捕获，那么就会闭合该作用域内的所有局部变量对应的upvalue
func (self *funcInfo) getJmpArgA() int {

	hasCapturedLocVars := false			// 是否有局部变量被闭包捕获
	minSlotOfLocVars := self.maxRegs	// 记录该作用域内局部变量的最小slot值

	for _, locVar := range self.locNames {
		if locVar.scopeLv == self.scopeLv {
			for v := locVar; v != nil && v.scopeLv == self.scopeLv; v = v.prev {  // 遍历在同一个scopeLv的局部变量
				if v.captured {		// 如果该局部变量被闭包捕获，那么就需要jmp指令去close闭包
					hasCapturedLocVars = true
				}

				if v.slot < minSlotOfLocVars && v.name[0] != '(' {  // for语句的自动生成的局部变量的变量名开头为'('
					minSlotOfLocVars = v.slot   // 该作用域内变量的最小slot值
				}
			}
		}
	}

	if hasCapturedLocVars {
		return minSlotOfLocVars + 1		// jmp指令计算需要close的闭包的最小slot - 1
	} else {
		return 0		// 0表示没有局部变量被捕获，即不需要进行close upvalue
	}
}


// 返回已经生成的最后一条指令的程序计数器(pc)，也就是该指令的索引
func (self *funcInfo) pc() int {

	return len(self.insts) - 1
}


// 修改指令iAsBx指令sBx操作数
func (self *funcInfo) fixSbx(pc, sBx int) {

	i := self.insts[pc]
	i = i << 18 >> 18			// 清除sBx操作数
	i = i | uint32(sBx+MAXARG_sBx)<<14
	self.insts[pc] = i
}


// 生成4中编码的模式

func (self *funcInfo) emitABC(opcode, a, b, c int) {

	i := b<<23 | c<<14 | a<<6 | opcode

	self.insts = append(self.insts, uint32(i))
}


func (self *funcInfo) emitABx(opcode, a, bx int)  {

	i := bx<<14 | a<<6 | opcode

	self.insts = append(self.insts, uint32(i))
}


func (self *funcInfo) emitAsBx(opcode, a, b int) {

	i := (b+MAXARG_sBx)<<14 | a<<6 | opcode

	self.insts = append(self.insts, uint32(i))
}


func (self *funcInfo) emitAx(opcode, ax int) {

	i := ax<<6 | opcode

	self.insts = append(self.insts, uint32(i))
}


// r[a] = r[b]
func (self *funcInfo) emitMove(a, b int) {
	self.emitABC(OP_MOVE, a, b, 0)
}


// r[a], r[a+1], ..., r[a+b] = nil
func (self *funcInfo) emitLoadNil(a, n int) {
	self.emitABC(OP_LOADNIL, a, n-1, 0)		// -1的原因是从R(a)开始赋值为nil
}


// r[a] = (bool)b; if (c) pc++
func (self *funcInfo) emitLoadBool(a, b, c int) {
	self.emitABC(OP_LOADBOOL, a, b, c)
}


// r[a] = kst[bx]
func (self *funcInfo) emitLoadK(a int, k interface{}) {
	idx := self.indexOfConstant(k)
	if idx < (1 << 18) {
		self.emitABx(OP_LOADK, a, idx)
	} else {
		self.emitABx(OP_LOADKX, a, 0)
		self.emitAx(OP_EXTRAARG, idx)
	}
}


// r[a], r[a+1], ..., r[a+b-2] = vararg
func (self *funcInfo) emitVararg(a, n int) {
	self.emitABC(OP_VARARG, a, n+1, 0)
}


// r[a] = emitClosure(proto[bx])
func (self *funcInfo) emitClosure(a, bx int) {
	self.emitABx(OP_CLOSURE, a, bx)
}


// r[a] = {}
func (self *funcInfo) emitNewTable(a, nArr, nRec int) {
	self.emitABC(OP_NEWTABLE,
		a, Int2fb(nArr), Int2fb(nRec))
}


// r[a][(c-1)*FPF+i] := r[a+i], 1 <= i <= b
func (self *funcInfo) emitSetList(a, b, c int) {
	self.emitABC(OP_SETLIST, a, b, c)
}


// r[a][(c-1)*FPF+i] := r[a+i], 1 <= i <= b
func (self *funcInfo) emitGetTable(a, b, c int) {
	self.emitABC(OP_GETTABLE, a, b, c)
}


// r[a][rk(b)] = rk(c)
func (self *funcInfo) emitSetTable(a, b, c int) {
	self.emitABC(OP_SETTABLE, a, b, c)
}


// r[a] = upval[b]
func (self *funcInfo) emitGetUpval(a, b int) {
	self.emitABC(OP_GETUPVAL, a, b, 0)
}


// upval[b] = r[a]
func (self *funcInfo) emitSetUpval(a, b int) {
	self.emitABC(OP_SETUPVAL, a, b, 0)
}


// r[a] = upval[b][rk(c)]
func (self *funcInfo) emitGetTabUp(a, b, c int) {
	self.emitABC(OP_GETTABUP, a, b, c)
}


// upval[a][rk(b)] = rk(c)
func (self *funcInfo) emitSetTabUp(a, b, c int) {
	self.emitABC(OP_SETTABUP, a, b, c)
}


// r[a], ..., r[a+c-2] = r[a](r[a+1], ..., r[a+b-1])
func (self *funcInfo) emitCall(a, nArgs, nRet int) {
	self.emitABC(OP_CALL, a, nArgs+1, nRet+1)
}


// return r[a](r[a+1], ... ,r[a+b-1])
func (self *funcInfo) emitTailCall(a, nArgs int) {
	self.emitABC(OP_TAILCALL, a, nArgs+1, 0)
}


// return r[a], ... ,r[a+b-2]
func (self *funcInfo) emitReturn(a, n int) {
	self.emitABC(OP_RETURN, a, n+1, 0)
}


// r[a+1] := r[b]; r[a] := r[b][rk(c)]
func (self *funcInfo) emitSelf(a, b, c int) {
	self.emitABC(OP_SELF, a, b, c)
}


// pc+=sBx; if (a) close all upvalues >= r[a - 1]
func (self *funcInfo) emitJmp(a, sBx int) int {
	self.emitAsBx(OP_JMP, a, sBx)
	return len(self.insts) - 1
}


// if not (r[a] <=> c) then pc++
func (self *funcInfo) emitTest(a, c int) {
	self.emitABC(OP_TEST, a, 0, c)
}


// if (r[b] <=> c) then r[a] := r[b] else pc++
func (self *funcInfo) emitTestSet(a, b, c int) {
	self.emitABC(OP_TESTSET, a, b, c)
}


func (self *funcInfo) emitForPrep(a, sBx int) int {
	self.emitAsBx(OP_FORPREP, a, sBx)
	return len(self.insts) - 1
}


func (self *funcInfo) emitForLoop(a, sBx int) int {
	self.emitAsBx(OP_FORLOOP, a, sBx)
	return len(self.insts) - 1
}


func (self *funcInfo) emitTForCall(a, c int) {
	self.emitABC(OP_TFORCALL, a, 0, c)
}


func (self *funcInfo) emitTForLoop(a, sBx int) {
	self.emitAsBx(OP_TFORLOOP, a, sBx)
}


// r[a] = op r[b]
func (self *funcInfo) emitUnaryOp(op, a, b int) {
	switch op {
	case TOKEN_OP_NOT:
		self.emitABC(OP_NOT, a, b, 0)
	case TOKEN_OP_BNOT:
		self.emitABC(OP_BNOT, a, b, 0)
	case TOKEN_OP_LEN:
		self.emitABC(OP_LEN, a, b, 0)
	case TOKEN_OP_UNM:
		self.emitABC(OP_UNM, a, b, 0)
	}
}


// r[a] = rk[b] op rk[c]
// arith & bitwise & relational
func (self *funcInfo) emitBinaryOp(op, a, b, c int) {
	if opcode, found := arithAndBitwiseBinops[op]; found {
		self.emitABC(opcode, a, b, c)
	} else {
		// 比较运算符需要添加额外的指令
		switch op {
		case TOKEN_OP_EQ:
			self.emitABC(OP_EQ, 1, b, c)
		case TOKEN_OP_NE:
			self.emitABC(OP_EQ, 0, b, c)
		case TOKEN_OP_LT:
			self.emitABC(OP_LT, 1, b, c)
		case TOKEN_OP_GT:
			self.emitABC(OP_LT, 1, c, b)
		case TOKEN_OP_LE:
			self.emitABC(OP_LE, 1, b, c)
		case TOKEN_OP_GE:
			self.emitABC(OP_LE, 1, c, b)
		}
		self.emitJmp(0, 1)
		self.emitLoadBool(a, 0, 1)
		self.emitLoadBool(a, 1, 0)
	}
}